{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.7 Softmax回归的简洁实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch.nn import init\n",
    "import numpy as np\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data_fashion_mnist(batch_size):\n",
    "    #batch_size = 32    # 256\n",
    "    if sys.platform.startswith('win'):\n",
    "        num_workers = 0\n",
    "    else:\n",
    "        num_workers = 4\n",
    "    mnist_train = torchvision.datasets.FashionMNIST(\n",
    "        root='./Datasets', train=True, download=False, \n",
    "        transform=torchvision.transforms.ToTensor())\n",
    "    mnist_test = torchvision.datasets.FashionMNIST(\n",
    "        root='./Datasets', train=False, download=False,\n",
    "        transform=torchvision.transforms.ToTensor()) \n",
    "    train_iter = torch.utils.data.DataLoader(mnist_train, \n",
    "                batch_size=batch_size, shuffle=True, num_workers=num_workers)\n",
    "    test_iter = torch.utils.data.DataLoader(mnist_test, \n",
    "                batch_size=batch_size, shuffle=False, num_workers=num_workers)\n",
    "    return train_iter, test_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32    # 256\n",
    "train_iter, test_iter = load_data_fashion_mnist(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_inputs = 784\n",
    "num_outputs = 10\n",
    "\n",
    "class LinearNet(nn.Module):\n",
    "    def __init__(self, num_inputs, num_outputs):\n",
    "        super(LinearNet, self).__init__()\n",
    "        self.linear = nn.Linear(num_inputs, num_outputs)\n",
    "        \n",
    "    def forward(self, x):    # x shape: (batch, 1, 28, 28)\n",
    "        y = self.linear(x.view(x.shape[0], -1))\n",
    "        return y\n",
    "\n",
    "net = LinearNet(num_inputs, num_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlattenLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(FlattenLayer, self).__init__()\n",
    "        \n",
    "    def forward(self, x):    # x shape: (batch, *, *, ...)\n",
    "        return x.view(x.shape[0], -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "net = nn.Sequential(\n",
    "    OrderedDict([\n",
    "        ('flatten', FlattenLayer()),\n",
    "        ('linear', nn.Linear(num_inputs, num_outputs))\n",
    "    ])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "init.normal_(net.linear.weight, mean=0, std=0.01)\n",
    "init.constant_(net.linear.bias, val=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD(net.parameters(), lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_accuracy(data_iter, net):\n",
    "    acc_sum, n = 0.0, 0\n",
    "    for X, y in data_iter:\n",
    "        acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()\n",
    "        n += y.shape[0]\n",
    "    return acc_sum / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, \n",
    "             params=None, lr=None, optimizer=None):\n",
    "    for epoch in range(num_epochs):\n",
    "        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0\n",
    "        for X, y in train_iter:\n",
    "            y_hat = net(X);\n",
    "            l = loss(y_hat, y).sum()\n",
    "            \n",
    "            # 梯度清零\n",
    "            if optimizer is not None:\n",
    "                optimizer.zero_grad()\n",
    "            elif params is not None and params[0].grad is not None:\n",
    "                for param in params:\n",
    "                    param.grad.data.zero_()\n",
    "            \n",
    "            # 反向求loss函数对各变量的偏导数，即梯度\n",
    "            l.backward()\n",
    "            if optimizer is None:\n",
    "                sgd(params, lr, batch_size)\n",
    "            else:\n",
    "                optimizer.step()\n",
    "                \n",
    "            train_l_sum += l.item()\n",
    "            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()\n",
    "            n += y.shape[0]\n",
    "            \n",
    "        test_acc = evaluate_accuracy(test_iter, net)\n",
    "        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f' \n",
    "             % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 5\n",
    "train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size,\n",
    "         None, None, optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
